class ImageGenerator:
    """图片生成模型工具类"""

    def __init__(self, config: AIConfig):
        self.config = config

    def generate_with_dalle(self, prompt: str, size: str = "1024x1024", quality: str = "standard"):
        """使用DALL-E生成图片"""
        import openai

        client = openai.OpenAI(api_key=self.config.configs["openai"]["api_key"])

        response = client.images.generate(
            model="dall-e-3",
            prompt=prompt,
            size=size,
            quality=quality,
            n=1,
        )

        return response.data[0].url

    def generate_with_stable_diffusion(self, prompt: str, **kwargs):
        """使用Stable Diffusion生成图片"""
        api_key = self.config.configs["stability"]["api_key"]

        response = requests.post(
            f"{self.config.configs['stability']['base_url']}/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            },
            json={
                "text_prompts": [{"text": prompt}],
                "cfg_scale": 7,
                "height": 1024,
                "width": 1024,
                "samples": 1,
                "steps": 30,
                **kwargs
            }
        )

        if response.status_code == 200:
            data = response.json()
            return data["artifacts"][0]["base64"]
        else:
            raise Exception(f"生成失败: {response.text}")

    def edit_image(self, image_path: str, prompt: str):
        """图片编辑"""
        # 实现图片编辑逻辑
        pass

    def upscale_image(self, image_path: str, scale: int = 2):
        """图片超分辨率"""
        # 实现图片放大逻辑
        pass